# Copyright (c) 2023, Tri Dao.

# To run the huggingface implementation, we first need to convert the weights:
# https://github.com/huggingface/transformers/pull/21955
# python -m transformers.models.llama.convert_llama_weights_to_hf --input_dir $CHECKPOINT_DIR/llama --model_size 7B --output_dir $CHECKPOINT_DIR$/llama/7B-hf
# and repeat for 13B, 30B, 65B

import os
import time
from pathlib import Path
current_dir = Path(__file__).parent.absolute()

import torch

# from transformers import LlamaConfig, LlamaTokenizer
# from transformers.models.llama.modeling_llama import LlamaForCausalLM

from flash_attn.models.gpt import GPTLMHeadModel, combine_state_dicts_tp
from flash_attn.models.llama import remap_state_dict_meta_llama, llama_config_to_gpt2_config
from flash_attn.models.llama import config_from_checkpoint, state_dicts_from_checkpoint
from flash_attn.utils.pretrained import state_dict_from_pretrained
from flash_attn.utils.generation import update_graph_cache
from flagai.data.tokenizer import Tokenizer
from GPT_mp_tools import shard_state_dict_tp
import bminf

## TODO
# checkpoints = "/share/project/64node-bmt-flashatten/state_dict/"
# model_name = "llama-30b-en"
# print('*'*20, "model_name", model_name, flush=True)
# cache_dir = checkpoints + model_name
# #print('*'*20, "cache_dir", cache_dir)
# tokenizer = Tokenizer.from_pretrained(model_name, cache_dir=cache_dir)
#print('*'*20, "tokenizer", tokenizer)
# config_file = cache_dir + "/config.json"
# # avoid sync loading models in case of Mem OOM
# if env_args.bmt_async_load:
#     import time
#     time.sleep(2*60*(trainer.local_rank%8))

# from flagai.model.llama_model import LLAMAModel
# model = LLAMAModel.init_from_json(config_file=config_file)

tokenizer = Tokenizer.from_pretrained("llama-30b-en", 
                                      cache_dir="./gpt2_new_100k/")

model_name='30B'
device = "cuda:0"

## /share/project/64node-bmt-flashatten/checkpoints/Aquila-30b-64n8g-from-scratch/7500
checkpoint_path = '/share/projset/baaishare/baai-mrnd/llama/llama'
config = llama_config_to_gpt2_config(config_from_checkpoint(checkpoint_path, model_name))
config.vocab_size=100008
config.use_cache = True
config.attn_pdrop = 0.0
config.resid_pdrop = 0.0
config.fused_bias_fc = True
config.use_flash_attn = True
config.fused_mlp = False  # We don't have fused GatedMLP yet
config.fused_dropout_add_ln = True
config.residual_in_fp32 = True


print(config)
world_size = 4
dtype = torch.float16

if not torch.distributed.is_initialized():
        torch.distributed.init_process_group(backend='nccl', init_method='env://')
device = f'cuda:{torch.distributed.get_rank()}'
# assert world_size <= torch.distributed.get_world_size()
# Need this, otherwise when we capture the graph the process for GPU 1 would run on both
# GPU0 and GPU1 and things would hang
torch.cuda.set_device(device)

from apex.transformer import parallel_state
parallel_state.initialize_model_parallel(tensor_model_parallel_size_=world_size)
rank = parallel_state.get_tensor_model_parallel_rank()
process_group = parallel_state.get_tensor_model_parallel_group()

print(f"rank is {rank}")

input_ids = tokenizer.encode_plus("请问你是谁？")["input_ids"]
print(input_ids)

input_ids = torch.tensor(input_ids)[None, ].to(device)
# print(config)
print(f"input shape is {input_ids.shape}")


# if not rotary, we load the weight from HF but ignore the position embeddings.
# The model would be nonsense but it doesn't matter for the test.
# model = GPTLMHeadModel.from_pretrained("./aquila-30b-7500/llama-30b-en/", config, 
#                                         strict=True, device=device,
#                                         dtype=dtype, process_group=process_group,
#                                         world_size=world_size, rank=rank)
# model.eval()


model = GPTLMHeadModel(config, 
                       device=device, 
                       process_group=process_group,
                       dtype=dtype)

# with torch.cuda.device(0):
        # model = bminf.wrapper(model, quantization=False, memory_limit=20 << 40)

# sd = torch.load("/share/project/64node-bmt-flashatten/checkpoints/Aquila-30b-64n8g-from-scratch/13500/pytorch_model.bin", map_location="cpu")["module"]
sd = torch.load("./transe_fc3_pytorch_model.bin", map_location="cpu")

sd = shard_state_dict_tp(sd, config, world_size=world_size, rank=rank)

model.load_state_dict(sd, strict=True)

print(f"参数加载成功")

# out = model.generate(input_ids=input_ids, max_length=max_length,
#                          eos_token_id=eos_token_id, fused_ft_kernel=True,
#                          return_dict_in_generate=True, output_scores=True, timing=True,
#                          teacher_outputs=out_hf.sequences)

out = model.generate(input_ids=input_ids, max_length=256,
                         vocab_size=config.vocab_size, fused_ft_kernel=True,
                         return_dict_in_generate=True, output_scores=True, timing=True)

print(out.sequences)

# def test_llama_generation(model_name):
#     checkpoint_path = Path(os.environ.get('CHECKPOINT_DIR',
#                                           current_dir.parent.parent / 'checkpoints')) / 'llama'

#     dtype = torch.float16
#     device = 'cuda'
#     config = llama_config_to_gpt2_config(config_from_checkpoint(checkpoint_path, model_name))
#     config.use_flash_attn = True
#     config.fused_bias_fc = True
#     config.fused_mlp = False  # We don't have fused GatedMLP yet
#     config.fused_dropout_add_ln = True
#     config.residual_in_fp32 = True

#     tokenizer = LlamaTokenizer.from_pretrained(Path(checkpoint_path) / f'{model_name}-hf')
#     eos_token_id = tokenizer.eos_token_id

#     torch.manual_seed(0)
#     batch_size = 1
#     seqlen = 100
#     max_length = 150
#     input_ids = torch.randint(0, config.vocab_size, (batch_size, seqlen), dtype=torch.long,
#                               device=device)

#     model_hf = LlamaForCausalLM.from_pretrained(Path(checkpoint_path) / f'{model_name}-hf',
#                                                 torch_dtype=dtype, device_map={"": device})
#     model_hf.eval()
#     print("HF fp16")
#     torch.cuda.synchronize()
#     start = time.time()
#     out_hf = model_hf.generate(input_ids=input_ids, max_length=max_length,
#                                return_dict_in_generate=True, output_scores=True)
#     torch.cuda.synchronize()
#     print(f'Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms')
#     del model_hf

#     model_ref = LlamaForCausalLM.from_pretrained(Path(checkpoint_path) / f'{model_name}-hf',
#                                                  device_map={"": device})
#     model_ref.eval()
#     with torch.no_grad():
#         logits_ref = model_ref(out_hf.sequences).logits[:, (seqlen - 1):-1]
#     del model_ref

#     ckpt_state_dicts = state_dicts_from_checkpoint(checkpoint_path, model_name)
#     pretrained_state_dicts = [remap_state_dict_meta_llama(s, config) for s in ckpt_state_dicts]
#     pretrained_state_dict = combine_state_dicts_tp(pretrained_state_dicts, config)
#     model = GPTLMHeadModel(config, device=device, dtype=dtype)
#     model.load_state_dict(pretrained_state_dict, strict=False)
#     model.eval()

#     print('Without CUDA graph')
#     torch.cuda.synchronize()
#     start = time.time()
#     out = model.generate(input_ids=input_ids, max_length=max_length,
#                          eos_token_id=eos_token_id, fused_ft_kernel=True,
#                          return_dict_in_generate=True, output_scores=True, timing=True,
#                          teacher_outputs=out_hf.sequences)
#     torch.cuda.synchronize()
#     print(f'Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms')

#     # Capture graph outside the timing loop
#     batch_size, seqlen_og = input_ids.shape
#     model._decoding_cache = update_graph_cache(model, None, batch_size, seqlen_og, max_length)
#     print('With CUDA graph')
#     torch.cuda.synchronize()
#     start = time.time()
#     out_cg = model.generate(input_ids=input_ids, max_length=max_length,
#                             fused_ft_kernel=True, cg=True,
#                             return_dict_in_generate=True, output_scores=True, timing=True,
#                             teacher_outputs=out_hf.sequences)
#     torch.cuda.synchronize()
#     print(f'Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms')

#     with torch.no_grad():
#         logits_parallel = model(out_hf.sequences).logits[:, (seqlen - 1):-1]
#     logits_hf = torch.stack(out_hf.scores, dim=1)
#     logits = torch.stack(out.scores, dim=1)
#     logits_cg = torch.stack(out_cg.scores, dim=1)

#     del model

#     hf_error = (logits_hf - logits_ref).abs().max().item()
#     # For some reason logits_parallel is off by quite a bit more than 2x
#     assert (logits_parallel - logits_ref).abs().max().item() < 8 * hf_error

#     print(f'HF fp16 logits max diff: {hf_error}')
#     print(f'Logits max diff: {(logits - logits_parallel).abs().max().item() }')
#     assert (logits - logits_parallel).abs().max().item() < 2 * hf_error
#     print(f'Logits CG max diff: {(logits_cg - logits_parallel).abs().max().item() }')
#     assert torch.equal(logits_cg, logits)